package org.csu.onlineedubackend.config;

import cn.dev33.satoken.stp.StpUtil;
import cn.dev33.satoken.exception.NotLoginException;
import org.csu.onlineedubackend.websocket.ChatWebSocketHandler;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.http.server.ServerHttpRequest;
import org.springframework.http.server.ServerHttpResponse;
import org.springframework.http.server.ServletServerHttpRequest;
import org.springframework.web.socket.WebSocketHandler;
import org.springframework.web.socket.config.annotation.EnableWebSocket;
import org.springframework.web.socket.config.annotation.WebSocketConfigurer;
import org.springframework.web.socket.config.annotation.WebSocketHandlerRegistry;
import org.springframework.web.socket.server.HandshakeInterceptor;
import org.springframework.web.socket.server.standard.ServletServerContainerFactoryBean;

import jakarta.servlet.http.HttpServletRequest;
import java.util.Enumeration;
import java.util.Map;

@Configuration
@EnableWebSocket
public class WebSocketConfig implements WebSocketConfigurer {

    @Autowired
    private ChatWebSocketHandler chatWebSocketHandler;

    @Override
    public void registerWebSocketHandlers(WebSocketHandlerRegistry registry) {
        registry.addHandler(chatWebSocketHandler, "/api/ws/chat")
                .addInterceptors(new ChatHandshakeInterceptor())
                .setAllowedOrigins("*"); // 允许所有来源
    }

    @Bean
    public ServletServerContainerFactoryBean createWebSocketContainer() {
        ServletServerContainerFactoryBean container = new ServletServerContainerFactoryBean();
        container.setMaxTextMessageBufferSize(8192);
        container.setMaxBinaryMessageBufferSize(8192);
        return container;
    }

    // 握手拦截器，用于在WebSocket连接建立时提取用户信息和验证token
    private class ChatHandshakeInterceptor implements HandshakeInterceptor {
        @Override
        public boolean beforeHandshake(ServerHttpRequest request, ServerHttpResponse response, 
                                      WebSocketHandler wsHandler, Map<String, Object> attributes) throws Exception {
            if (request instanceof ServletServerHttpRequest) {
                HttpServletRequest servletRequest = ((ServletServerHttpRequest) request).getServletRequest();
        
                // 打印所有请求头，用于调试
                System.out.println("WebSocket握手请求头信息:");
                Enumeration<String> headerNames = servletRequest.getHeaderNames();
                while (headerNames.hasMoreElements()) {
                    String headerName = headerNames.nextElement();
                    System.out.println(headerName + ": " + servletRequest.getHeader(headerName));
                }
                
                // 从请求参数中获取token
                String token = servletRequest.getParameter("token");
                System.out.println("WebSocket握手收到token参数: " + token);
                
                // 尝试从请求头中获取token
                String headerToken = servletRequest.getHeader("satoken");
                if (headerToken != null && !headerToken.isEmpty()) {
                    System.out.println("WebSocket握手从请求头获取到token: " + headerToken);
                    // 优先使用请求头中的token
                    token = headerToken;
                }
                
                if (token == null || token.isEmpty()) {
                    System.out.println("WebSocket握手失败：token为空");
                    return false;
                }
                
                try {
                    // 使用Sa-Token验证token并获取用户ID
                    Object loginId = StpUtil.getLoginIdByToken(token);
                    if (loginId == null) {
                        System.out.println("WebSocket握手失败：token无效，无法获取用户ID");
                        return false;
                    }
                    
                    // 将用户ID放入WebSocket会话属性
                    attributes.put("userId", Integer.parseInt(loginId.toString()));
                    System.out.println("WebSocket握手成功：用户ID=" + loginId);
                    return true;
                } catch (Exception e) {
                    System.out.println("WebSocket握手失败：token验证异常 - " + e.getMessage());
                    e.printStackTrace();
                    return false;
                }
            }
            
            System.out.println("WebSocket握手失败：请求不是ServletServerHttpRequest类型");
            return false;
        }

        @Override
        public void afterHandshake(ServerHttpRequest request, ServerHttpResponse response, 
                                  WebSocketHandler wsHandler, Exception exception) {
            if (exception != null) {
                System.out.println("WebSocket握手完成，但有异常：" + exception.getMessage());
            } else {
                System.out.println("WebSocket握手完成，连接建立成功");
            }
        }
    }
}